#include "precompiled.h"
#include <math.h>
#include <gsl/gsl_fft_complex.h>
#include <gsl/gsl_complex_math.h>
#include "Config.h"
#include "Util.h"
#include "Environment.h"
#include "NatcarSimViz.h"

struct FFTData {
	gsl_complex_packed_array data;	// array of double
	int N;
	GLfloat r, g, b;
	double max;

	FFTData(gsl_complex_packed_array _data, int _N, GLfloat _r, GLfloat _g, GLfloat _b) :
		data(_data),	
		N(_N),
		r(_r),
		g(_g),
		b(_b)
	{
		max = 0;
		// get the maximum magnitude
		for(int i = 0; i < N; i++) {
			double x = gsl_complex_abs(*(gsl_complex *)&data[2*i]);
			if(x > max) {
				max = x;
			}
		}
	}

	~FFTData() {
		delete[] data;
	}
};

NatcarSimViz::NatcarSimViz(Environment &env) : 
	OpenGLWindow(), 
	FieldSizeX(100),
	FieldSizeY(100),
	FieldZ(0),
	FieldScale(1.0e8),
	WindowSplit(0.5),
	PlaybackSpeed(1.0),
	mResult(env.GetResult()),
	mEnv(env),
	mPause(false),
	mStep(0),
	mDrawField(false),
	mFieldTex(0),
	mDrawFFT(false),
	mFFTInitialized(false)
{
	typedef json::grammar<char>::array jarray;
	typedef json::grammar<char>::object jobject;

	jobject &o = Config::Get("Viz");

	Config::TryAssign(FieldSizeX, "FieldSizeX", o);
	Config::TryAssign(FieldSizeX, "FieldSizeX", o);
	Config::TryAssign(FieldSizeY, "FieldSizeY", o);
	Config::TryAssign(FieldZ, "FieldZ", o);
	Config::TryAssign(FieldScale, "FieldScale", o);
	Config::TryAssign(WindowSplit, "WindowSplit", o);
	Config::TryAssign(PlaybackSpeed, "PlaybackSpeed", o);

	// set up the plots
	jarray &jplots = boost::any_cast<jarray>(*o.find("Plots")->second);
	mPlots.resize(jplots.size());
	unsigned int i = 0;
	for(jarray::const_iterator it = jplots.begin(); it != jplots.end(); it++, i++) {
		Plot &p = mPlots[i];
		p.id = i;
		jobject &o1 = boost::any_cast<jobject>(**it);
		Config::TryAssign(p.SecsPerDiv, "SecsPerDiv", o1);
		Config::TryAssign(p.HDivs, "HDivs", o1);
		Config::TryAssign(p.VDivs, "VDivs", o1);
	}

	// go through the signals and save the properties
	jobject &a = boost::any_cast<jobject>(*o.find("Signals")->second);
	mSignalProps.resize(a.size());
	i = 0;
	for(jobject::const_iterator it = a.begin(); i < a.size(); it++, i++)
	{
		jobject &s = boost::any_cast<jobject>(*it->second);
		SignalProps &c = mSignalProps[i];

		mSignalMap[it->first] = i;

		jarray &rgb = boost::any_cast<jarray>(*s.find("Rgb")->second);
		Config::AssignArray(c.Rgb, rgb);

		c.Name = it->first;
		Config::TryAssign(c.Scale, "Scale", s);
		Config::TryAssign(c.Plot, "Plot", s);
		Config::TryAssign(c.Display, "Display", s);
	}
}

NatcarSimViz::~NatcarSimViz() {
	DestroyFFT();
	DestroyMagneticField();
}

void NatcarSimViz::Run(HINSTANCE hInstance) {
	// initialize GDI
	Gdiplus::GdiplusStartupInput gdiplusStartupInput;
	ULONG_PTR gdiplusToken;
	Gdiplus::GdiplusStartup(&gdiplusToken, &gdiplusStartupInput, NULL);

	// initialize window and application
	InitializeWindow(hInstance, "Natcar Simulation Vizualizer");
	Util::LoadExtensions();
	InitGL();

	mResult.Reset();
	// now show the window and draw the image. Want a nominal framerate of 60fps
	ShowWindow(17);

	Gdiplus::GdiplusShutdown(gdiplusToken);
}

void NatcarSimViz::InitGL() {
	glClearColor(0, 0, 0, 1.0);			// Black Background
	glShadeModel(GL_FLAT);
	//glEnable(GL_BLEND);
	//glBlendFunc (GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);	// want to allow alpha blending
	//glDisable(GL_DEPTH_TEST);
	glMatrixMode(GL_MODELVIEW);
	glLoadIdentity();
}

void NatcarSimViz::OnResize(GLsizei width, GLsizei height) {
	// calculate the sizes of the sub-viewports
	mWindowWidth = width;
	mWindowHeight = height;

	// set up viewport 1 (left side of screen)
	// have to maintain the aspect ratio of the world
	mVp1[0] = 0;
	mVp1[1] = 0;
	double vp1Width = WindowSplit * width;
	double aspect = vp1Width / height;
	double aspectWorld = mEnv.Width / mEnv.Height;
	if(aspect >= aspectWorld) {
		// if the screen is wider than the world, fit to height
		mVp1[2] = (GLsizei)(height * aspectWorld);
		mVp1[3] = (GLsizei)height;
	} else {
		mVp1[2] = (GLsizei)vp1Width;
		mVp1[3] = (GLsizei)(vp1Width / aspectWorld);
	}
	
	// divide up the right half of the screen into however many plots we have
	int plotHeight = height / mPlots.size();
	for(int i = 0; i < mPlots.size(); i++) {
		Plot &p = mPlots[i];
		p.vp[0] = (int)(vp1Width + 1);
		p.vp[1] = i * plotHeight;
		p.vp[2] = (GLsizei)(width - p.vp[0]);
		p.vp[3] = plotHeight;
	}
}

void NatcarSimViz::OnUpdate() {
	if(mDrawFFT) {
		// nothing
	} else if(mPause) {
		// do nothing
	} else {
		int index = (int)(0.017 * PlaybackSpeed * mStep / mEnv.TimeStep);
		if(index >= mResult.GetSize()) {
			index = 0;
			mStep = 0;
		}
		mResult.Seek(index);
		mStep++;
	}
}

// put drawing logic here
void NatcarSimViz::OnDisplay() {
	glClear(GL_COLOR_BUFFER_BIT);
	if(mDrawFFT) {
		DrawFFT();
	} else {
		DrawViewport1();
		DrawViewport2();
	}
}

void NatcarSimViz::OnKeyDown(char c) {
	switch(c) {
	case 'F':	// toggle DFT visualization
		mDrawFFT = !mDrawFFT;
		std::cerr << (mDrawFFT?"Entering FFT mode. Hit 'f' to exit":"Exiting FFT mode") << "\n";
		InitFFT();
		break;
	case 'M':	// toggle magnetic field visualization
		mDrawField = !mDrawField;
		std::cerr << (mDrawField?"Enabling Field Rendering":"Disabling Field Rendering") << "\n";
		InitMagneticField();
		break;
	case 'P':	// pause the simulation
		mPause = !mPause;
		std::cerr << (mPause?"Simulation paused. Hit 'p' to unpause":"Simulation unpaused") << "\n";
		break;
	}
}

// Draw the simulation visualization viewport (left side)
void NatcarSimViz::DrawViewport1() {
	glViewport(mVp1[0], mVp1[1], mVp1[2], mVp1[3]);
	glMatrixMode(GL_PROJECTION);					// Select The Projection Matrix
	glLoadIdentity();								// Reset The Projection Matrix
	glOrtho(0, mEnv.Width, 0, mEnv.Height, -1.0, 1.0);
	glMatrixMode(GL_MODELVIEW);
	glLoadIdentity();
	
	glColor3f(0, 0, 1.0);
	glBegin(GL_LINE_STRIP);
		glVertex2f(0, 0);
		glVertex2f(mEnv.Width, 0);
		glVertex2f(mEnv.Width, mEnv.Height);
		glVertex2f(0, mEnv.Height);
		glVertex2f(0, 0);
	glEnd();

	if(mDrawField) {
		DrawMagneticField();
	}
	// draw the track in white
	glColor3f(1.0, 1.0, 1.0);
	mEnv.GetTrack().Draw();

	const SimResult::Data *data = mResult.GetData();
	glColor3f(227.0/255.0, 108.0/255.0, 10.0/255.0);
	// draw the car's trail
	glBegin(GL_LINE_STRIP);
	for(int i = 0; i < mResult.GetIndex(); i++) {
		glVertex2f(data[i].carPos[0], data[i].carPos[1]);
	}
	glEnd();

	// draw the position of the car at this instant as a rectangle
	const Robot &robot = mEnv.GetRobot();
	const ublas::vector<double> &pos = mResult.Current().carPos;
	double direction = mResult.Current().carDirection;
	// translate by Width / 2
	glColor3f(1.0, 0, 0);
	glTranslated(pos[0], pos[1], 0);
	glRotatef((GLfloat)(direction * 180.0 / M_PI), 0, 0, 1.0);
	glRectd(0, -(GLfloat)robot.GetWidth()/2, (GLfloat)robot.GetLength(), (GLfloat)robot.GetWidth()/2);
}

void NatcarSimViz::DrawAxis(const Plot &p) {
	// draw the divs in light grey
	glColor3f(0.8, 0.8, 0.8);
	glBegin(GL_LINES);
	// first draw horizontal divs
	for(int i = 0; i < p.HDivs; i++) {
		glVertex2f(-1.0, (i - p.HDivs/2) * 2.0/p.HDivs);
		glVertex2f(1.0, (i - p.HDivs/2) * 2.0/p.HDivs);
	}
	// then draw vertical divs
	for(int i = 0; i < p.VDivs; i++) {
		glVertex2f((i - p.VDivs/2) * 2.0/p.VDivs, 1.0);
		glVertex2f((i - p.VDivs/2) * 2.0/p.VDivs, -1.0);
	}
	glEnd();

	glColor3f(0.0, 0.0, 0.0);
	glBegin(GL_LINES);
		// x-axis (horizontal)
		glVertex2f(-1.0, 0);
		glVertex2f(1.0, 0);
		// y-axis (vertical)
		glVertex2f(0, -1.0);
		glVertex2f(0, 1.0);

		// upper and lower bounds
		glVertex2f(-1.0, 1.0);
		glVertex2f(1.0, 1.0);
		glVertex2f(-1.0, -1.0);
		glVertex2f(1.0, -1.0);
	glEnd();
}

// Draw the information viewport (right side)
// Subdivide into two viewports: top and bottom
void NatcarSimViz::DrawViewport2() {
	for(int i = 0; i < mPlots.size(); i++) {
		DrawPlot(mPlots[i]);
	}
}

void NatcarSimViz::DrawPlot(const Plot &p) {
	// set up viewport
	glViewport(p.vp[0], p.vp[1], p.vp[2], p.vp[3]);
	glMatrixMode(GL_PROJECTION);					// Select The Projection Matrix
	glLoadIdentity();								// Reset The Projection Matrix
	glOrtho(-1.0, 1.0, -1.0, 1.0, -1.0, 1.0);
	glMatrixMode(GL_MODELVIEW);
	glLoadIdentity();

	// draw white background
	glColor3f(1.0, 1.0, 1.0);
	glRectd(-1.0, -1.0, 2.0, 2.0);

	// draw the axis in black
	glColor3f(0, 0, 0);
	DrawAxis(p);

	// draw all the signals for this plot
	DrawSignals(p);
}

void NatcarSimViz::DrawSignals(const Plot &p) {
	// based on SecsPerDiv, how many points of the signal should be plotted?
	int numPoints = p.VDivs * p.SecsPerDiv / mEnv.TimeStep;

	// draw all the signals that have requested to be displayed on this plot
	for(int i = 0; i < mSignalProps.size(); i++) {
		const SignalProps &s = mSignalProps[i];
		if(s.Plot != p.id || !s.Display)
			continue;
		glColor3f(s.Rgb[0] / 255.0, s.Rgb[1] / 255.0, s.Rgb[2] / 255.0);
		// look up this signal's index in the SimResult
		int j = mResult.GetSignalIndex(s.Name);
		DrawSignal(&mResult.GetData()->signals[j], sizeof(SimResult::Data), mResult.GetSize(), mResult.GetIndex(), numPoints, s.Scale);
	}
}

// stride is the number of bytes between elements.
// For example, if data is a dense array of doubles, then stride should be sizeof(double)
void NatcarSimViz::DrawSignal(const void *data, int stride, int length, int index, int numPoints, double scale) {
	int start = index - numPoints/2;
	glBegin(GL_LINE_STRIP);
	for(int i = 0; i < numPoints && (i + start < length); i++) {
		if(start + i >= 0) {
			glVertex2f(2.0*i/numPoints - 1.0, scale * *(const double *)((const char *)data + stride*(i + start)));
		}
	}
	glEnd();
}

// Compute the magnetic field in the specified resolution
// and save the result in a texture so that we can draw it quickly
void NatcarSimViz::InitMagneticField() {
	if(mFieldTex != 0) {
		return;
	}
	std::cout << "Computing Magnetic Field\n";

	// calculate the magnetic field over the entire environment in the specified resolution
	float *buf = new float[FieldSizeX * FieldSizeY];
	double xscale = mEnv.Width / FieldSizeX;
	double yscale = mEnv.Height / FieldSizeY;
	int lastPercent = 0;
	for(int i = 0; i < FieldSizeY; i++) {
		for(int j = 0; j < FieldSizeX; j++) {
			ublas::vector<double> p(3);
			p[0] = j * xscale;
			p[1] = i * yscale;
			p[2] = FieldZ;
			vec field = mEnv.GetTrack().GetFieldAt(p);
			//double mag = ublas::norm_2(field);
			int index = i * FieldSizeX + j;
			buf[index] = (float)(abs(field[2]) * FieldScale);
		}
		//std::cerr << "Done with row " << i << "\n";
		// show our progress by printing dots accross the screen
		int percent = 80 * i / FieldSizeY;
		int numPeriods = percent - lastPercent;
		for(int i = 0; i < numPeriods; i++) {
			std::cout << ".";
		}
		lastPercent = percent;
	}

	GLuint tex;
	glGenTextures(1, &tex);
	glBindTexture(GL_TEXTURE_2D, tex);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
	glTexEnvf(GL_TEXTURE_ENV, GL_TEXTURE_ENV_MODE, GL_DECAL);
	glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, FieldSizeX, FieldSizeY, 0, GL_GREEN, GL_FLOAT, buf);
	delete[] buf;

	mFieldTex = tex;

	std::cout << "\nDone.\n";
}

void NatcarSimViz::DestroyMagneticField() {
	if(mFieldTex != 0) {
		glDeleteTextures(1, &mFieldTex);
		mFieldTex = 0;
	}
}

void NatcarSimViz::DrawMagneticField() {
	glEnable(GL_TEXTURE_2D);
	glBindTexture(GL_TEXTURE_2D, mFieldTex);
	glBegin (GL_QUADS);
		glTexCoord2f (0.0, 0.0);
		glVertex2f (0.0, 0.0);
		glTexCoord2f (1.0, 0.0);
		glVertex2f ((GLfloat)mEnv.Width, 0.0);
		glTexCoord2f (1.0, 1.0);
		glVertex2f ((GLfloat)mEnv.Width, (GLfloat)mEnv.Height);
		glTexCoord2f (0.0, 1.0);
		glVertex2f (0.0, (GLfloat)mEnv.Height);
	glEnd ();
	glDisable(GL_TEXTURE_2D);
}

void NatcarSimViz::PrintGLInfo() {
	const GLubyte *xtns = glGetString(GL_EXTENSIONS);
	std::cerr << "OpenGL Extensions:\n\n" << xtns << "\n\n";

	// get number of bits in alpha bitplane
	int alphaBits;
	glGetIntegerv(GL_DEPTH_BITS, &alphaBits);
	std::cerr << "alphaBits: " << alphaBits << "\n";
}



void NatcarSimViz::InitFFT() {
	if(mFFTInitialized) {
		return;
	}
	std::cerr << "Computing FFT for lsense, rsense, and error ... ";
	/*
	// allocate memory for the buffers
	int L = mResult.GetSize();
	int N = Util::NextHighestPowerOf2((unsigned int)L);
	gsl_complex_packed_array lsenseDFT = new double[2*N];
	gsl_complex_packed_array rsenseDFT = new double[2*N];
	gsl_complex_packed_array errorDFT = new double[2*N];

	// zero out everything
	memset(lsenseDFT, 0, sizeof(*lsenseDFT)*2*N);
	memset(rsenseDFT, 0, sizeof(*rsenseDFT)*2*N);
	memset(errorDFT, 0, sizeof(*errorDFT)*2*N);

	// copy data into each buffer
	const SimResult::Data *data = mResult.GetData();
	for(int i = 0; i < L; i++) {
		lsenseDFT[2*i] = data[i].lSense;
		rsenseDFT[2*i] = data[i].rSense;
		errorDFT[2*i] = data[i].error;
	}

	gsl_fft_complex_radix2_forward(lsenseDFT, 1, N);
	gsl_fft_complex_radix2_forward(rsenseDFT, 1, N);
	gsl_fft_complex_radix2_forward(errorDFT, 1, N);

	// assign atomically
	mLsenseDFT = new FFTData(lsenseDFT, N, 1.0, 0, 0);
	mRsenseDFT = new FFTData(rsenseDFT, N, 0, 0, 1.0);
	mErrorDFT = new FFTData(errorDFT, N, 0, 1.0, 0);
	mFFTInitialized = true;
	*/
	std::cerr << "Done.\n";
}

void NatcarSimViz::DestroyFFT() {
	if(!mFFTInitialized) {
		return;
	}
	/*
	delete mLsenseDFT;
	delete mRsenseDFT;
	delete mErrorDFT;
	*/
	mFFTInitialized = false;
}

void NatcarSimViz::DrawFFT() {
	/*
	// set up viewport
	glViewport(0, 0, mWindowWidth, mWindowHeight);
	glMatrixMode(GL_PROJECTION);					// Select The Projection Matrix
	glLoadIdentity();								// Reset The Projection Matrix
	glOrtho(-0.01, 1.0, -0.01, 1.0, -1.0, 1.0);
	glMatrixMode(GL_MODELVIEW);
	glLoadIdentity();

	// draw white background
	glColor3f(1.0, 1.0, 1.0);
	glRectf(-0.1, -0.1, 1.1, 1.1);

	// draw the axis in black
	glColor3f(0, 0, 0);
	glBegin(GL_LINES);
		// x-axis (horizontal)
		glVertex2f(-1.0, 0);
		glVertex2f(1.0, 0);
		// y-axis (vertical)
		glVertex2f(0, -1.0);
		glVertex2f(0, 1.0);
	glEnd();
	
	DrawFFTMag(*mLsenseDFT);
	DrawFFTMag(*mRsenseDFT);
	//DrawFFTMag(*mErrorDFT);
	*/
}

// draw only from 0 to pi
void NatcarSimViz::DrawFFTMag(const FFTData &dft) {
	glColor3f(dft.r, dft.g, dft.b);
	glBegin(GL_LINE_STRIP);
	for(int i = 0; i < dft.N/2; i++) {
		glVertex2f((double)i/(dft.N/2-1), gsl_complex_abs(*(gsl_complex *)&dft.data[2*i]) / dft.max);
	}
	glEnd();
}


